{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7c092a3b",
   "metadata": {},
   "source": [
    "In this tutorial, we'll demonstrate the Argoverse 2.0 map API, and visualize some of the map data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f25e1e66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from argparse import Namespace\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d70c7396",
   "metadata": {},
   "outputs": [],
   "source": [
    "# path to where the logs live\n",
    "dataroot = \"/Users/jlambert/Downloads/av2-v7-2022-01-15\"\n",
    "\n",
    "# unique log identifier\n",
    "log_id = \"adcf7d18-0510-35b0-a2fa-b4cea13a6d76\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f919b29b",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Namespace(**{\"dataroot\": Path(dataroot), \"log_id\": Path(log_id)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15a8196d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "from pathlib import Path\n",
    "from typing import List\n",
    "\n",
    "import matplotlib\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "import av2.geometry.polyline_utils as polyline_utils\n",
    "import av2.rendering.vector as vector_plotting_utils\n",
    "from av2.datasets.sensor.av2_sensor_dataloader import AV2SensorDataLoader\n",
    "from av2.map.map_api import ArgoverseStaticMap, LaneSegment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88d1aaf7",
   "metadata": {},
   "source": [
    "First, we'll plot the lane graph and crosswalks, with crosswalks colored purple, and the lanes colored green."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a02b3b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# scaled to [0,1] for matplotlib.\n",
    "PURPLE_RGB = [201, 71, 245]\n",
    "PURPLE_RGB_MPL = np.array(PURPLE_RGB) / 255\n",
    "\n",
    "DARK_GRAY_RGB = [40, 39, 38]\n",
    "DARK_GRAY_RGB_MPL = np.array(DARK_GRAY_RGB) / 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d5ae748",
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_log_teaser(args: argparse.Namespace) -> None:\n",
    "    \"\"\"For a single log, render all local crosswalks in green, and pedestrian crossings in purple, in a bird's eye view.\"\"\"\n",
    "    log_map_dirpath = Path(args.dataroot) / args.log_id / \"map\"\n",
    "    avm = ArgoverseStaticMap.from_map_dir(log_map_dirpath, build_raster=False)\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 10))\n",
    "    ax = fig.add_subplot()\n",
    "\n",
    "    for _, ls in avm.vector_lane_segments.items():\n",
    "        # right_ln_bnd\n",
    "        # left_ln_bnd\n",
    "        vector_plotting_utils.draw_polygon_mpl(\n",
    "            ax, ls.polygon_boundary, color=\"g\", linewidth=0.5\n",
    "        )\n",
    "        vector_plotting_utils.plot_polygon_patch_mpl(\n",
    "            ls.polygon_boundary, ax, color=\"g\", alpha=0.2\n",
    "        )\n",
    "\n",
    "    # plot all pedestrian crossings\n",
    "    for _, pc in avm.vector_pedestrian_crossings.items():\n",
    "        vector_plotting_utils.draw_polygon_mpl(ax, pc.polygon, color=\"m\", linewidth=0.5)\n",
    "        vector_plotting_utils.plot_polygon_patch_mpl(\n",
    "            pc.polygon, ax, color=\"m\", alpha=0.2\n",
    "        )\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e27fe09",
   "metadata": {},
   "outputs": [],
   "source": [
    "single_log_teaser(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56a98bd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_raster_layers(args: argparse.Namespace) -> None:\n",
    "    \"\"\"Visualize the ground surface height/elevation map, w/ a colorbar indicating the value range.\n",
    "\n",
    "    Also, visualize side-by-side plots of the 3 raster arrays -- ground height, drivable area, ROI.\n",
    "    \"\"\"\n",
    "    log_map_dirpath = Path(args.dataroot) / args.log_id / \"map\"\n",
    "    avm = ArgoverseStaticMap.from_map_dir(log_map_dirpath, build_raster=True)\n",
    "\n",
    "    height_array = avm.raster_ground_height_layer.array\n",
    "    ax = plt.subplot()\n",
    "    plt.title(\"Ground surface height (@ 30 centimeter resolution).\")\n",
    "    img = plt.imshow(np.flipud(height_array))\n",
    "\n",
    "    divider = make_axes_locatable(ax)\n",
    "    cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "    plt.colorbar(img, cax=cax)\n",
    "    plt.show()\n",
    "\n",
    "    plt.figure(figsize=(20, 10))\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.imshow(np.flipud(height_array))\n",
    "    plt.title(\"Ground Surface Height\")\n",
    "\n",
    "    plt.subplot(1, 3, 2)\n",
    "    da_array = avm.raster_drivable_area_layer.array\n",
    "    plt.imshow(np.flipud(da_array))\n",
    "    plt.title(\"Drivable Area (rasterized \\nfrom vector polygons)\")\n",
    "\n",
    "    plt.subplot(1, 3, 3)\n",
    "    roi_array = avm.raster_roi_layer.array\n",
    "    plt.imshow(np.flipud(roi_array))\n",
    "    plt.title(\"Region of Interest (ROI)\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bdc3926",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_raster_layers(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f45e82da",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_lane_segments(\n",
    "    ax: matplotlib.axes.Axes,\n",
    "    lane_segments: List[LaneSegment],\n",
    "    lane_color: np.ndarray = DARK_GRAY_RGB_MPL,\n",
    ") -> None:\n",
    "    \"\"\"Plot lane segments.\"\"\"\n",
    "    for ls in lane_segments:\n",
    "        pts_city = ls.polygon_boundary\n",
    "        ALPHA = 1.0  # 0.1\n",
    "        vector_plotting_utils.plot_polygon_patch_mpl(\n",
    "            polygon_pts=pts_city, ax=ax, color=lane_color, alpha=ALPHA, zorder=1\n",
    "        )\n",
    "\n",
    "        for bound_type, bound_city in zip(\n",
    "            [ls.left_mark_type, ls.right_mark_type],\n",
    "            [ls.left_lane_boundary, ls.right_lane_boundary],\n",
    "        ):\n",
    "            if \"YELLOW\" in bound_type:\n",
    "                mark_color = \"y\"\n",
    "            elif \"WHITE\" in bound_type:\n",
    "                mark_color = \"w\"\n",
    "            else:\n",
    "                mark_color = \"grey\"  # \"b\" lane_color #\n",
    "\n",
    "            LOOSELY_DASHED = (0, (5, 10))\n",
    "\n",
    "            if \"DASHED\" in bound_type:\n",
    "                linestyle = LOOSELY_DASHED\n",
    "            else:\n",
    "                linestyle = \"solid\"\n",
    "\n",
    "            if \"DOUBLE\" in bound_type:\n",
    "                left, right = polyline_utils.get_double_polylines(\n",
    "                    polyline=bound_city.xyz[:, :2], width_scaling_factor=0.1\n",
    "                )\n",
    "                ax.plot(\n",
    "                    left[:, 0],\n",
    "                    left[:, 1],\n",
    "                    mark_color,\n",
    "                    alpha=ALPHA,\n",
    "                    linestyle=linestyle,\n",
    "                    zorder=2,\n",
    "                )\n",
    "                ax.plot(\n",
    "                    right[:, 0],\n",
    "                    right[:, 1],\n",
    "                    mark_color,\n",
    "                    alpha=ALPHA,\n",
    "                    linestyle=linestyle,\n",
    "                    zorder=2,\n",
    "                )\n",
    "            else:\n",
    "                ax.plot(\n",
    "                    bound_city.xyz[:, 0],\n",
    "                    bound_city.xyz[:, 1],\n",
    "                    mark_color,\n",
    "                    alpha=ALPHA,\n",
    "                    linestyle=linestyle,\n",
    "                    zorder=2,\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f7dec27",
   "metadata": {},
   "outputs": [],
   "source": [
    "def argoverse2_sensor_dataset_teaser(\n",
    "    args: argparse.Namespace, save_plot: bool = False\n",
    ") -> None:\n",
    "    \"\"\"Visualize both ego-vehicle poses and the per-log local vector map.\n",
    "\n",
    "    Crosswalks are plotted in purple. Lane segments plotted in dark gray. Ego-pose in red.\n",
    "    \"\"\"\n",
    "    loader = AV2SensorDataLoader(data_dir=args.dataroot, labels_dir=args.dataroot)\n",
    "\n",
    "    fig = plt.figure(1, figsize=(10, 10))\n",
    "    ax = fig.add_subplot(111)\n",
    "\n",
    "    log_map_dirpath = Path(args.dataroot) / args.log_id / \"map\"\n",
    "\n",
    "    avm = ArgoverseStaticMap.from_map_dir(log_map_dirpath, build_raster=False)\n",
    "\n",
    "    # retain every pose first.\n",
    "    traj_ns = loader.get_subsampled_ego_trajectory(args.log_id, sample_rate_hz=1e9)\n",
    "    # now, sample @ 1 Hz\n",
    "    traj_1hz = loader.get_subsampled_ego_trajectory(args.log_id, sample_rate_hz=1.0)\n",
    "    med_x, med_y = np.median(traj_ns, axis=0)\n",
    "\n",
    "    # Derive plot area from trajectory (with radius defined in infinity norm).\n",
    "    view_radius_m = 50\n",
    "    xlims = [med_x - view_radius_m, med_x + view_radius_m]\n",
    "    ylims = [med_y - view_radius_m, med_y + view_radius_m]\n",
    "\n",
    "    crosswalk_color = PURPLE_RGB_MPL\n",
    "    CROSSWALK_ALPHA = 0.6\n",
    "    for pc in avm.get_scenario_ped_crossings():\n",
    "        vector_plotting_utils.plot_polygon_patch_mpl(\n",
    "            polygon_pts=pc.polygon[:, :2],\n",
    "            ax=ax,\n",
    "            color=crosswalk_color,\n",
    "            alpha=CROSSWALK_ALPHA,\n",
    "            zorder=3,\n",
    "        )\n",
    "\n",
    "    plot_lane_segments(ax=ax, lane_segments=avm.get_scenario_lane_segments())\n",
    "\n",
    "    ax.plot(traj_ns[:, 0], traj_ns[:, 1], color=\"r\", zorder=4, label=\"Ego-vehicle pose\")\n",
    "    ax.scatter(\n",
    "        traj_1hz[:, 0], traj_1hz[:, 1], 100, facecolors=\"none\", edgecolors=\"r\", zorder=4\n",
    "    )  # marker='o', color=\"r\")\n",
    "\n",
    "    plt.axis(\"equal\")\n",
    "    plt.xlim(*xlims)\n",
    "    plt.ylim(*ylims)\n",
    "    plt.title(f\"Log {args.log_id}\")\n",
    "    plt.axis(\"off\")\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac070630",
   "metadata": {},
   "outputs": [],
   "source": [
    "argoverse2_sensor_dataset_teaser(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b3cc09c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
